# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
import random
from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig


@GENERATOR_REGISTRY.register("ascend_generate_grouped_bias_add_grad_v2")
class GroupedBiasAddGradV2Generator(CaseGenerator):

    def __init__(self, config):
        super().__init__(config)

    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        grad_y_shape = case_config.inputs[0].shape
        if len(grad_y_shape) == 3:
            case_config.inputs[1].shape = [0]
        else:
            case_config.inputs[1].shape = [grad_y_shape[0]]
            case_config.inputs[1].dtype = random.choice(["int32", "int64"])
        return case_config